import os

import numpy as np
import cv2 as cv
import argparse
from visualizer import OccupancyVisualizer
import open3d as o3d
import colorsys
import mmengine
from nuscenes.nuscenes import NuScenes

occ3d_colors_map = np.array(
    [
        [0, 0, 0],          # 其他 (others)               黑色
        [255, 120, 50],     # 障碍物 (barrier)            橙色
        [255, 192, 203],    # 自行车 (bicycle)            粉色
        [255, 255, 0],      # 公交车 (bus)                黄色
        [0, 150, 245],      # 小汽车 (car)                蓝色
        [0, 255, 255],      # 工程车 (construction_vehicle) 青色
        [255, 127, 0],      # 摩托车 (motorcycle)         深橙色
        [255, 0, 0],        # 行人 (pedestrian)           红色
        [255, 240, 150],    # 交通锥 (traffic_cone)       浅黄色
        [135, 60, 0],       # 拖车 (trailer)              棕色
        [160, 32, 240],     # 卡车 (truck)                紫色
        [255, 0, 255],      # 可行驶路面 (driveable_surface) 深粉色
        [139, 137, 137],    # 其他平面 (other_flat)        深灰色
        [75, 0, 75],        # 人行道 (sidewalk)           深紫色
        [150, 240, 80],     # 地形 (terrain)              浅绿色
        [230, 230, 250],    # 建筑物和墙 (manmade)         浅蓝白色
        [0, 175, 0],        # 植被 (vegetation)           绿色
        [255, 255, 255],    # 自由空间 (free)             白色
    ]
)
# change to BGR
occ3d_colors_map = occ3d_colors_map[:, ::-1]

openocc_colors_map = np.array(
    [
        [0, 150, 245],      # car                  blue         √
        [160, 32, 240],     # truck                purple       √
        [135, 60, 0],       # trailer              brown        √
        [255, 255, 0],      # bus                  yellow       √
        [0, 255, 255],      # construction_vehicle cyan         √
        [255, 192, 203],    # bicycle              pink         √
        [255, 127, 0],      # motorcycle           dark orange  √
        [255, 0, 0],        # pedestrian           red          √
        [255, 240, 150],    # traffic_cone         light yellow
        [255, 120, 50],     # barrier              orange
        [255, 0, 255],      # driveable_surface    dark pink
        [139, 137, 137],    # other_flat           dark red
        [75, 0, 75],        # sidewalk             dard purple
        [150, 240, 80],     # terrain              light green
        [230, 230, 250],    # manmade              white
        [0, 175, 0],        # vegetation           green
        [255, 255, 255],    # Free                 White
    ]
)

foreground_idx = [0, 1, 2, 3, 4, 5, 6, 7]


def parse_args():
    parse = argparse.ArgumentParser('')
    parse.add_argument('--inputs', type=str, required=True, help='path of the inputs directory or file')
    parse.add_argument('--ignore_classes', type=int, nargs='+', default=[15, 17], 
                      help='要忽略的类别序号，例如：15 17 表示忽略建筑物和自由空间\n'
                           '0: 其他\n'
                           '1: 障碍物\n'
                           '2: 自行车\n'
                           '3: 公交车\n'
                           '4: 小汽车\n'
                           '5: 工程车\n'
                           '6: 摩托车\n'
                           '7: 行人\n'
                           '8: 交通锥\n'
                           '9: 拖车\n'
                           '10: 卡车\n'
                           '11: 可行驶路面\n'
                           '12: 其他平面\n'
                           '13: 人行道\n'
                           '14: 地形\n'
                           '15: 建筑物和墙\n'
                           '16: 植被\n'
                           '17: 自由空间')
    args = parse.parse_args()
    return args

def arange_according_to_scene(infos, nusc, vis_scene):
    scenes = dict()

    for i, info in enumerate(infos):
        scene_token = nusc.get('sample', info['token'])['scene_token']
        scene_meta = nusc.get('scene', scene_token)
        scene_name = scene_meta['name']
        if not scene_name in scenes:
            scenes[scene_name] = [info]
        else:
            scenes[scene_name].append(info)

    vis_scenes = dict()
    if len(vis_scene) == 0:
        vis_scenes = scenes
    else:
        for scene_name in vis_scene:
            vis_scenes[scene_name] = scenes[scene_name]

    return vis_scenes

def vis_occ_scene_on_3d(vis_scenes_infos,
                        vis_scene,
                        vis_path,
                        pred_path,
                        dataset_type='occ3d',
                        load_camera_mask=False,
                        voxel_size=(0.4, 0.4, 0.4),
                        vis_gt=True,
                        vis_flow=False,
                        car_model=None,
                        background_color=(255, 255, 255),
                        line_width=0.0,
                        ignore_classes=None,
                        save_video=True,
                        video_format='avi',  # 默认使用 AVI 格式
                        fps=30,
                        ):
    # define free_cls
    free_cls = 16 if dataset_type == 'occ3d' else 17

    # 处理要忽略的类别
    if ignore_classes is None:
        ignore_classes = []
    
    # 获取类别名称列表（中英对照）
    class_names = {
        'others': '其他',
        'barrier': '障碍物',
        'bicycle': '自行车',
        'bus': '公交车',
        'car': '小汽车',
        'construction_vehicle': '工程车',
        'motorcycle': '摩托车',
        'pedestrian': '行人',
        'traffic_cone': '交通锥',
        'trailer': '拖车',
        'truck': '卡车',
        'driveable_surface': '可行驶路面',
        'other_flat': '其他平面',
        'sidewalk': '人行道',
        'terrain': '地形',
        'manmade': '建筑物和墙',
        'vegetation': '植被',
        'free': '自由空间'
    }
    
    # 将类别名称转换为索引
    ignore_indices = []
    for cls_name in ignore_classes:
        if cls_name in class_names:
            idx = list(class_names.keys()).index(cls_name)
            if idx not in ignore_indices:
                ignore_indices.append(idx)
    
    # 添加255（未知类别）到忽略列表
    if 255 not in ignore_indices:
        ignore_indices.append(255)

    print("\n=== 忽略的类别信息 ===")
    for idx in sorted(ignore_indices):
        if idx < len(class_names):
            cls_name = list(class_names.keys())[idx]
            print(f"类别 {idx}: {class_names[cls_name]} ({cls_name})")
        elif idx == 255:
            print("类别 255: 未知类别 (unknown)")
    print("====================\n")

    # load car model
    if car_model is not None:
        car_model_mesh = o3d.io.read_triangle_mesh(car_model)
        angle = np.pi / 2  # 90 度
        R = car_model_mesh.get_rotation_matrix_from_axis_angle(np.array([angle, 0, 0]))
        car_model_mesh.rotate(R, center=car_model_mesh.get_center())
        car_model_mesh.scale(0.25, center=car_model_mesh.get_center())
        current_center = car_model_mesh.get_center()
        new_center = np.array([0, 0, 0.5])
        translation = new_center - current_center
        car_model_mesh.translate(translation)
        car_model_mesh.compute_vertex_normals()
    else:
        car_model_mesh = None

    # check vis path
    mmengine.mkdir_or_exist(vis_path)
    
    # 处理输入路径
    if os.path.isdir(pred_path):
        # 如果是目录，直接使用目录下的所有场景
        scene_dirs = [d for d in os.listdir(pred_path) if os.path.isdir(os.path.join(pred_path, d))]
        vis_scene = scene_dirs if not vis_scene else vis_scene
    else:
        # 如果是单个文件，只处理该文件
        scene_dirs = [os.path.basename(os.path.dirname(pred_path))]
        vis_scene = scene_dirs

    # 创建可视化器实例（在循环外创建，所有场景共用）
    occ_visualizer = OccupancyVisualizer(
        color_map=occ3d_colors_map if dataset_type == 'occ3d' else openocc_colors_map,
        background_color=background_color
    )

    # 只在第一帧时加载视角设置
    view_json = None
    if os.path.exists('view.json'):
        view_json = o3d.io.read_pinhole_camera_parameters('view.json')

    for scene_name in vis_scene:
        print(f"\r处理场景: {scene_name}", end='', flush=True)
        scene_infos = vis_scenes_infos[scene_name]
        buffer_vis_path = '{}/{}'.format(vis_path, scene_name)
        mmengine.mkdir_or_exist(buffer_vis_path)

        current_frame = 0

        def on_next_frame():
            nonlocal current_frame, view_json
            if current_frame < len(scene_infos):
                info = scene_infos[current_frame]
                save_path = os.path.join(buffer_vis_path, str(current_frame))
                
                # 可视化场景数据
                if vis_gt:
                    # 可视化GT
                    occ_path = info['occ_path']
                    if dataset_type == 'openocc':
                        occ_path = occ_path.replace('gts', 'openocc_v2')
                    occ_label_path = os.path.join(occ_path, 'labels.npz')
                    occ_label = np.load(occ_label_path)
                    occ_semantics = occ_label['semantics']

                    if load_camera_mask:
                        assert 'mask_camera' in occ_label.keys()
                        mask_camera = occ_label['mask_camera']
                        occ_semantics[mask_camera == 0] = 255
                    
                    # 检查是否有flow数据
                    if vis_flow and 'flow' in occ_label.keys():
                        occ_flow = occ_label['flow']
                    else:
                        occ_flow = None
                else:
                    # 可视化预测结果
                    token = info['token']
                    if os.path.isdir(pred_path):
                        occ_label_path = os.path.join(pred_path, scene_name, token, 'pred.npz')
                    else:
                        occ_label_path = pred_path
                        
                    occ_label = np.load(occ_label_path)
                    occ_semantics = occ_label['pred']
                    
                    # 检查是否有flow数据
                    if vis_flow:
                        if 'flow' in occ_label.keys():
                            occ_flow = occ_label['flow']
                        elif 'flows' in occ_label.keys():
                            occ_flow = occ_label['flows']
                        else:
                            occ_flow = None
                    else:
                        occ_flow = None

                # 执行可视化
                occ_visualizer.vis_occ(
                    occ_semantics,
                    occ_flow=occ_flow,
                    ignore_labels=ignore_indices,
                    voxelSize=voxel_size,
                    range=[-40.0, -40.0, -1.0, 40.0, 40.0, 5.4],
                    save_path=save_path,
                    wait_time=-1,
                    view_json=view_json if current_frame == 0 and scene_name == vis_scene[0] else None,
                    car_model_mesh=car_model_mesh,
                    line_width=line_width,
                )
                print(f"\r处理场景: {scene_name} - 帧 {current_frame+1}/{len(scene_infos)}", end='', flush=True)
                current_frame += 1
            else:
                raise StopIteration

        # 设置回调函数
        occ_visualizer.on_next_frame = on_next_frame

        # 显示第一帧
        if len(scene_infos) > 0:
            try:
                if not hasattr(occ_visualizer, 'o3d_vis') or occ_visualizer.o3d_vis is None:
                    occ_visualizer.o3d_vis = occ_visualizer._initialize_o3d_vis()
                on_next_frame()
                occ_visualizer.show()
            except StopIteration:
                print(f"\r场景 {scene_name} 播放完毕", end='', flush=True)
                # 生成视频
                if save_video:
                    print(f"\r正在生成视频...", end='', flush=True)
                    vis_occ_semantics = []
                    frame_count = 0
                    
                    if not os.path.exists(buffer_vis_path):
                        print(f"\r错误：目录不存在: {buffer_vis_path}")
                        return
                    
                    # 收集所有帧
                    while True:
                        img_path = os.path.join(buffer_vis_path, f"{frame_count}.png")
                        if os.path.exists(img_path):
                            img = cv.imread(img_path)
                            if img is not None:
                                vis_occ_semantics.append(img)
                            # os.remove(img_path)
                            frame_count += 1
                        else:
                            break

                    if vis_occ_semantics:
                        # 创建视频保存目录
                        video_dir = os.path.join(vis_path, 'videos')
                        if vis_gt:
                            video_dir = os.path.join(video_dir, 'gt')
                        else:
                            video_dir = os.path.join(video_dir, 'pred')
                        if vis_flow:
                            video_dir = os.path.join(video_dir, 'flow')
                        
                        os.makedirs(video_dir, exist_ok=True)

                        # 保存视频
                        if video_format.lower() == 'avi':
                            fourcc = cv.VideoWriter_fourcc(*'MJPG')  # 使用 MJPG 编码器
                            video_ext = '.avi'
                        elif video_format.lower() == 'mp4':
                            try:
                                # 尝试使用 H.264
                                fourcc = cv.VideoWriter_fourcc(*'avc1')
                                video_ext = '.mp4'
                            except:
                                # 如果失败，回退到 MJPG
                                fourcc = cv.VideoWriter_fourcc(*'MJPG')
                                video_ext = '.avi'
                                print("\n警告：H.264 编码器不可用，使用 MJPG 编码器")
                        else:
                            raise ValueError(f"不支持的视频格式: {video_format}")

                        # 构建视频文件名
                        video_name = f"{scene_name}"
                        if vis_flow:
                            video_name += "_flow"
                        else:
                            video_name += "_occ"
                        video_name += video_ext
                        
                        video_path = os.path.join(video_dir, video_name)

                        # 获取图像尺寸
                        height, width = vis_occ_semantics[0].shape[:2]

                        try:
                            # 使用高质量设置创建视频写入器
                            video = cv.VideoWriter(video_path, fourcc, fps, (width, height))
                            if not video.isOpened():
                                print(f"\n错误：无法创建视频文件: {video_path}")
                                return

                            # 设置视频质量参数
                            try:
                                video.set(cv.VIDEOWRITER_PROP_QUALITY, 100)  # 最高质量
                            except:
                                print("\n警告：无法设置视频质量参数，使用默认设置")
                            
                            for img in vis_occ_semantics:
                                # 确保图像质量
                                if img.dtype != np.uint8:
                                    img = (img * 255).astype(np.uint8)
                                video.write(img)

                            video.release()
                            print(f"\n视频已保存: {video_path} ({len(vis_occ_semantics)}帧, {fps}fps, {width}x{height})")
                            
                        except Exception as e:
                            print(f"\n保存视频时发生错误: {str(e)}")
                            print("尝试使用备用编码器...")
                            try:
                                # 尝试使用 MJPG 编码器
                                backup_fourcc = cv.VideoWriter_fourcc(*'MJPG')
                                backup_path = video_path.replace('.mp4', '.avi')
                                video = cv.VideoWriter(backup_path, backup_fourcc, fps, (width, height))
                                if video.isOpened():
                                    for img in vis_occ_semantics:
                                        video.write(img)
                                    video.release()
                                    print(f"\n视频已使用备用编码器保存: {backup_path}")
                                else:
                                    print("\n错误：备用编码器也无法创建视频文件")
                            except Exception as e2:
                                print(f"\n备用编码器也失败: {str(e2)}")

                        # 保存GIF动图
                        try:
                            import imageio
                            gif_name = f"{scene_name}"
                            if vis_flow:
                                gif_name += "_flow"
                            else:
                                gif_name += "_occ"
                            gif_name += ".gif"
                            gif_path = os.path.join(video_dir, gif_name)
                            # imageio要求RGB顺序
                            gif_frames = [cv.cvtColor(img, cv.COLOR_BGR2RGB) for img in vis_occ_semantics]
                            imageio.mimsave(gif_path, 
                                            gif_frames, 
                                            duration=1.0/fps,
                                            loop=0) #无限循环
                            print(f"\nGIF已保存: {gif_path} ({len(gif_frames)}帧, {fps}fps, {width}x{height})")
                        except Exception as e:
                            print(f"\n保存GIF时发生错误: {str(e)}")
                    else:
                        print(f"\n警告：场景 {scene_name} 没有可用的帧")
                continue

    # 所有场景播放完毕后，关闭可视化窗口
    if hasattr(occ_visualizer, 'o3d_vis') and occ_visualizer.o3d_vis is not None:
        occ_visualizer.o3d_vis.destroy_window()
        occ_visualizer.o3d_vis.close()
        occ_visualizer._clear_o3d_vis()


def flow_to_color(vx, vy, max_magnitude=None):
    magnitude = np.sqrt(vx ** 2 + vy ** 2)
    angle = np.arctan2(vy, vx)

    hue = (angle + np.pi) / (2 * np.pi)

    if max_magnitude is None:
        max_magnitude = np.max(magnitude)

    saturation = np.clip(magnitude / max_magnitude, 0, 1)
    value = np.ones_like(saturation)

    hsv = np.stack((hue, saturation, value), axis=-1)
    rgb = np.apply_along_axis(lambda x: colorsys.hsv_to_rgb(*x), -1, hsv)
    rgb = (rgb * 255).astype(np.uint8)

    return rgb

def create_legend_circle(radius=1, resolution=500):
    x = np.linspace(-radius, radius, resolution)
    y = np.linspace(-radius, radius, resolution)
    X, Y = np.meshgrid(x, y)
    vx = X
    vy = Y
    magnitude = np.sqrt(vx ** 2 + vy ** 2)
    mask = magnitude <= radius

    vx = vx[mask]
    vy = vy[mask]

    colors = flow_to_color(vx, vy, max_magnitude=radius)

    legend_image = np.ones((resolution, resolution, 3), dtype=np.uint8) * 255
    legend_image[mask.reshape(resolution, resolution)] = colors

    return legend_image

if __name__ == '__main__':
    print('open3d version: {}, if you want to use viewcontrol, make sure using 0.16.0 version!!'.format(o3d.__version__))
    
    # 设置默认参数
    pkl_file = 'data/nuscenes/fbocc-nuscenes_infos_val.pkl'
    data_path = 'data/nuscenes'
    data_version = 'v1.0-mini'
    dataset_type = 'occ3d'
    vis_path = 'demo_out'
    car_model = 'tools/vis_tool/visualizer/3d_model.obj'
    line_width = 5
    inputs_path = '/home/fxp/Projects/FB-BEV/test/gsdoccdiss_r50_256x704_16f_8x4_24/Sat_Jun_21_02_49/occupancy_pred'
    ignore_classes = ['manmade', 'free']  # 默认忽略建筑物和自由空间
    
    # 视频保存参数
    save_video = True  # 是否保存视频
    video_format = 'avi'  # 默认使用 AVI 格式
    fps = 10  # 视频帧率
    
    # 检查vis path
    os.makedirs(vis_path, exist_ok=True)
    print(f"输出目录: {os.path.abspath(vis_path)}")
    print(f"目录权限: {os.access(vis_path, os.W_OK)}")
    print(f"可用空间: {os.statvfs(vis_path).f_bavail * os.statvfs(vis_path).f_frsize / 1024 / 1024:.2f}MB")
    
    # 加载数据
    pkl_data = mmengine.load(pkl_file)
    nusc = NuScenes(data_version, data_path)
    
    # 处理输入路径
    if os.path.isdir(inputs_path):
        vis_scene = [d for d in os.listdir(inputs_path) if os.path.isdir(os.path.join(inputs_path, d))]
    else:
        vis_scene = [os.path.basename(os.path.dirname(inputs_path))]
    
    vis_scenes_infos = arange_according_to_scene(pkl_data['infos'], nusc, vis_scene)
    
    vis_occ_scene_on_3d(
        vis_scenes_infos, 
        vis_scene, 
        vis_path, 
        inputs_path, 
        dataset_type=dataset_type, 
        vis_gt=False, 
        car_model=car_model,
        line_width=line_width,
        ignore_classes=ignore_classes,
        save_video=save_video,
        video_format=video_format,
        fps=fps,
    )
